import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.tree import DecisionTreeClassifier
from sklearn.linear_model import LogisticRegression
import joblib
from sklearn.metrics import roc_auc_score
from matplotlib import pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def data_format(data):

    # 绩效与离职情况
    x_0 = data[data['Attrition'] == 0].groupby(by='StandardHours')['Attrition'].count().reset_index()
    x_1 = data[data['Attrition'] == 1].groupby(by='StandardHours')['Attrition'].count().reset_index()
    print(x_0)
    print(x_1)
    x = np.arange(len(sorted(x_0['StandardHours'].unique())))
    print(x)
    width = 0.35  # 柱子宽度
    fig, ax = plt.subplots()
    bar0 = ax.bar(x - width / 2, x_0['Attrition'], 0.35, label='未离职')
    bar1 = ax.bar(x + width / 2, x_1['Attrition'], 0.35, label='已离职')
    plt.title('员工关系满意度与离职情况')
    plt.ylabel('人数')
    plt.xlabel('员工关系满意度')
    plt.legend()
    plt.xticks(x,x_0['StandardHours'])
    plt.tight_layout()
    # 显示数值标签
    for bars in [bar0, bar1]:
        for bar in bars:
            height = bar.get_height()
            ax.annotate(
                f"{height}",
                xy=(bar.get_x() + bar.get_width() / 2, height),
                xytext=(0, 3),  # 3 points vertical offset
                textcoords="offset points",
                ha="center",
                va="bottom"
            )


    # 员工关系满意度与离职情况
    x_0 = data[data['Attrition'] == 0].groupby(by='RelationshipSatisfaction')['Attrition'].count().reset_index()
    x_1 = data[data['Attrition'] == 1].groupby(by='RelationshipSatisfaction')['Attrition'].count().reset_index()
    print(x_0)
    print(x_1)
    x = np.arange(len(sorted(x_0['RelationshipSatisfaction'].unique())))
    print(x)
    width = 0.35  # 柱子宽度
    fig,ax = plt.subplots()
    bar0 = ax.bar(x - width/2, x_0['Attrition'], 0.35, label='未离职')
    bar1 = ax.bar(x + width/2, x_1['Attrition'], 0.35, label='已离职')
    plt.title('员工关系满意度与离职情况')
    plt.ylabel('人数')
    plt.xlabel('员工关系满意度')
    plt.legend()
    plt.xticks(x,['1分','2分','3分','4分'])
    plt.tight_layout()
    # 显示数值标签
    for bars in [bar0, bar1]:
        for bar in bars:
            height = bar.get_height()
            ax.annotate(
                f"{height}",
                xy=(bar.get_x() + bar.get_width() / 2, height),
                xytext=(0, 3),  # 3 points vertical offset
                textcoords="offset points",
                ha="center",
                va="bottom"
            )




    # 绩效与离职情况
    x_0 = data[data['Attrition'] == 0].groupby(by='PerformanceRating')['Attrition'].count().reset_index()
    x_1 = data[data['Attrition'] == 1].groupby(by='PerformanceRating')['Attrition'].count().reset_index()
    print(x_0)
    print(x_1)
    x = np.arange(len(sorted(x_0['PerformanceRating'].unique())))
    print(x)
    width = 0.35  # 柱子宽度
    fig,ax = plt.subplots()
    bar0 = ax.bar(x - width/2, x_0['Attrition'], 0.35, label='未离职')
    bar1 = ax.bar(x + width/2, x_1['Attrition'], 0.35, label='已离职')
    plt.title('绩效与离职情况')
    plt.ylabel('人数')
    plt.xlabel('绩效评分')
    plt.legend()
    plt.xticks(x,['3分','4分'])
    plt.tight_layout()
    # 显示数值标签
    for bars in [bar0, bar1]:
        for bar in bars:
            height = bar.get_height()
            ax.annotate(
                f"{height}",
                xy=(bar.get_x() + bar.get_width() / 2, height),
                xytext=(0, 3),  # 3 points vertical offset
                textcoords="offset points",
                ha="center",
                va="bottom"
            )



    # 涨薪百分比与离职情况
    x_0 = data[data['Attrition'] == 0].groupby(by='PercentSalaryHike')['Attrition'].count().reset_index()
    x_1 = data[data['Attrition'] == 1].groupby(by='PercentSalaryHike')['Attrition'].count().reset_index()
    print(x_0)
    print(x_1)
    x = np.arange(len(sorted(x_0['PercentSalaryHike'].unique())))
    print(x)
    width = 0.35  # 柱子宽度
    fig,ax = plt.subplots()
    bar0 = ax.bar(x - width/2, x_0['Attrition'], 0.35, label='未离职')
    bar1 = ax.bar(x + width/2, x_1['Attrition'], 0.35, label='已离职')
    plt.title('涨薪百分比与离职情况')
    plt.ylabel('人数')
    plt.xlabel('涨薪比例%')
    plt.legend()
    plt.xticks(x,['11%','12%','13%','14%','15%','16%','17%','18%','19%','20%','21%','22%','23%','24%','25%'])
    plt.tight_layout()
    # 显示数值标签
    for bars in [bar0, bar1]:
        for bar in bars:
            height = bar.get_height()
            ax.annotate(
                f"{height}",
                xy=(bar.get_x() + bar.get_width() / 2, height),
                xytext=(0, 3),  # 3 points vertical offset
                textcoords="offset points",
                ha="center",
                va="bottom"
            )
    # plt.savefig('data/fig/涨薪百分比与离职情况.PNG')
    # plt.show()


    # 加班与离职情况
    x_0 = data[data['Attrition'] == 0].groupby(by='OverTime')['Attrition'].count().reset_index()
    x_1 = data[data['Attrition'] == 1].groupby(by='OverTime')['Attrition'].count().reset_index()
    print(x_0)
    x = np.arange(len(sorted(x_0['OverTime'].unique())))
    print(x)
    width = 0.35  # 柱子宽度
    fig,ax = plt.subplots()
    ax.bar(x - width/2, x_0['Attrition'], 0.35, label='未离职')
    ax.bar(x + width/2, x_1['Attrition'], 0.35, label='已离职')
    plt.title('加班与离职情况')
    plt.ylabel('人数')
    plt.xlabel('离职情况')
    plt.legend()
    plt.xticks(x,['不加班','加班'])
    plt.tight_layout()
    # plt.savefig('data/fig/加班与离职情况.PNG')
    # plt.show()


    # 已离职和未离职任职过的企业数量的人数
    x_0 = data[data['Attrition'] == 0].groupby(by='NumCompaniesWorked')['Attrition'].count().reset_index()
    x_1 = data[data['Attrition'] == 1].groupby(by='NumCompaniesWorked')['Attrition'].count().reset_index()
    companies = sorted(x_0['NumCompaniesWorked'].unique())
    x = np.arange(len(companies))  # 生成 0, 1, 2, ... 作为 x 轴位置
    width = 0.35  # 柱子宽度
    fig,ax = plt.subplots()
    ax.bar(x - width/2, x_0['Attrition'], 0.35, label='未离职人数')
    ax.bar(x + width/2, x_1['Attrition'], 0.35, label='已离职人数')
    plt.title('已离职和未离职任职过的企业数量的人数')
    plt.ylabel('人数')
    plt.xlabel('任职过的企业数量')
    plt.legend()
    plt.xticks(x)
    plt.tight_layout()
    # plt.savefig('data/fig/已离职和未离职任职过的企业数量的人数.PNG')
    # plt.show()


    # 已离职与未离职的平均薪资
    x_0 = data[data['Attrition'] == 0]['MonthlyIncome'].mean()
    x_1 = data[data['Attrition'] == 1]['MonthlyIncome'].mean()
    plt.figure()
    plt.bar(x=['已离职','未离职'],height=[x_1,x_0])
    plt.title('已离职与未离职的平均薪资')
    plt.xlabel('离职情况')
    plt.ylabel('平均薪资')
    # plt.savefig('data/fig/已离职与未离职的平均薪资.PNG')
    # plt.show()
    pass


def fit_model(data):
    y = data['Attrition']
    x = pd.get_dummies(data.drop(['Attrition','Over18', 'RelationshipSatisfaction', 'StandardHours'],axis=1),drop_first=False)
    print(x)
    # x = data.iloc[:,2:]
    # x.drop(['Over18', 'Department', 'EducationField', 'JobRole', 'MaritalStatus', 'BusinessTravel', 'RelationshipSatisfaction', 'StandardHours'],axis=1,inplace=True)
    # x['OverTime'] = pd.get_dummies(data['OverTime'],drop_first=True)
    # x['Gender'] = pd.get_dummies(data['Gender'], drop_first=True)
    # x = pd.concat([x, pd.get_dummies(data[['BusinessTravel', 'Department', 'EducationField', 'JobRole', 'MaritalStatus']])], axis=1)
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3, random_state=2,stratify=y)
    estimator = RandomForestClassifier(class_weight='balanced',n_estimators=1500,random_state=2)
    estimator.fit(x_train, y_train)
    # estimator = DecisionTreeClassifier(random_state=2)
    # estimator.fit(x_train, y_train)
    predict = estimator.predict(x_test)
    print(f'AUC指标: {roc_auc_score(y_test, predict)}')
    # joblib.dump(estimator, r'./model/model.pkl')

def predict():
    joblib.load(r'./model/model.pkl')

if __name__ == '__main__':
    data = pd.read_csv('./data/train.csv')
    # data_format(data)
    # fit_model(data)
    # 3.特征工程
    x = pd.get_dummies(data.drop(['Attrition', 'Over18', 'StandardHours', 'RelationshipSatisfaction'], axis=1),
                       drop_first=False)
    y = data['Attrition']
    from sklearn.model_selection import train_test_split

    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.3, random_state=42, stratify=y
    )
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.model_selection import GridSearchCV

    model = RandomForestClassifier(class_weight='balanced', random_state=22)

    param_grid1 = {'n_estimators': [100, 200, 400, 800, 1000, 1200, 1500], }
    model_best = GridSearchCV(model, param_grid=param_grid1, cv=3, scoring="roc_auc", n_jobs=-1)
    model_best.fit(x_train, y_train)
    print(type(model_best.best_params_))

    param_grid2 = {'min_samples_split': range(30, 100)}
    model_best2 = GridSearchCV(model_best.best_estimator_, param_grid=param_grid2, cv=5, scoring="roc_auc", n_jobs=-1)
    model_best2.fit(x_train, y_train)

    param_grid3 = {'max_features': range(1, 30)}
    model_best3 = GridSearchCV(model_best.best_estimator_, param_grid=param_grid3, cv=5, scoring="roc_auc", n_jobs=-1)
    model_best3.fit(x_train, y_train)

    param_grid4 = {'max_depth': range(3, 25)}
    model_best4 = GridSearchCV(model_best.best_estimator_, param_grid=param_grid4, cv=5, scoring="roc_auc", n_jobs=-1)
    model_best4.fit(x_train, y_train)
    print(model_best4.best_estimator_)
    print(model_best4.best_params_, "  ", "得分：", model_best4.best_score_)
    y_predict = model_best4.predict(x_test)  # 进行预测
    y_pred_proba = model_best4.predict_proba(x_test)[:, 1]  # 进行预测

    from sklearn.metrics import roc_auc_score, roc_curve
    # 计算AUC
    auc = roc_auc_score(y_test, y_pred_proba)
    print(f"测试集AUC: {auc:.4f}")